import numpy as np
import cv2 as cv
import os
import matplotlib.pyplot as plt
from PIL import Image
from glob import glob


def compute_RGB_difference(prev, curr):
    ret = prev - curr
    np.clip(ret, 0, 255, ret)
    return ret
    # return np.abs(prev - curr)


def cal_for_diff(frames):
    flow = []
    prev = frames[0]

    for i, frame_curr in enumerate(frames[1:]):
        curr = frame_curr

        tmp_flow = compute_RGB_difference(prev, curr)
        flow.append(tmp_flow)
        prev = curr

    return flow


def video2frames(video_path):
    frames = []
    video = cv.VideoCapture(video_path)
    frame_num = video.get(7)
    print('video: {} has {:d} frames\n'.format(video_path, int(frame_num)))
    ret, frame = video.read()

    while ret and video.isOpened():
        frames.append(frame)
        ret, frame = video.read()

    print('Transfer video to frames successfully\n')

    return frames


def save_diff(video_flows, flow_path):
    save_path = os.path.join(flow_path, 'diff')
    if not os.path.exists(save_path):
        os.mkdir(os.path.join(save_path))

    for i, flow in enumerate(video_flows):
        cv.imwrite(os.path.join(save_path, 'diff_{:05d}.jpg'.format(i)), flow)

    return save_path


def extract_flow(video_path, flow_path):
    frames = video2frames(video_path)
    flow = cal_for_diff(frames)
    path = save_diff(flow, flow_path)
    print('Complete: {}\n'.format(flow_path))
    return path


def display(path):
    plt.ion()

    files = os.listdir(path)
    files.sort()
    for file in files:
        img = Image.open(os.path.join(path, file)).convert('P')
        plt.title(file)
        plt.imshow(img)
        plt.pause(0.1)

    plt.ioff()
    plt.show()
    return


def clean_jpg(path):
    imgs = glob(os.path.join(path, '*.jpg'))
    for img in imgs:
        os.remove(img)

    print('clean all images in {}\n'.format(path))


if __name__ == '__main__':
    video_path = "C:/Users/18873/Desktop/videos/v_Archery_g20_c06.avi"
    save_path = '../assets/'
    ret_path = extract_flow(video_path, save_path)
    display(ret_path)
    # clean_jpg(os.path.join(save_path, 'diff'))
